from flask import Flask, render_template, request, jsonify, g
from flask_cors import CORS
from datetime import datetime
import os
from models import db, Scenario, DatabaseConnection, QueryHistory, CalculationParameter
from database_connector import DatabaseConnector
import pandas as pd
import json

# 初始化Flask应用
app = Flask(__name__)
CORS(app)

# 配置数据库
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///instance/calculation_system.db'
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
app.config['SECRET_KEY'] = os.environ.get('SECRET_KEY', 'dev-key-for-development-only')

# 初始化数据库
db.init_app(app)

# 确保instance目录存在
os.makedirs(os.path.join(app.root_path, 'instance'), exist_ok=True)

# 数据库连接缓存
db_connections = {}

# 主页路由
@app.route('/')
def index():
    return render_template('index.html')

# 测算场景相关API
@app.route('/api/scenarios', methods=['GET'])
def get_scenarios():
    scenarios = Scenario.query.all()
    return jsonify([{
        'id': s.id,
        'name': s.name,
        'description': s.description,
        'type': s.type,
        'db_connection_id': s.db_connection_id,
        'created_at': s.created_at.isoformat(),
        'updated_at': s.updated_at.isoformat()
    } for s in scenarios])

@app.route('/api/scenarios', methods=['POST'])
def create_scenario():
    data = request.json
    scenario = Scenario(
        name=data['name'],
        description=data.get('description', ''),
        type=data['type'],
        db_connection_id=data.get('db_connection_id'),
        created_at=datetime.now(),
        updated_at=datetime.now()
    )
    db.session.add(scenario)
    db.session.commit()
    return jsonify({'id': scenario.id, 'message': 'Scenario created successfully'}), 201

@app.route('/api/scenarios/<int:scenario_id>', methods=['GET'])
def get_scenario(scenario_id):
    scenario = Scenario.query.get_or_404(scenario_id)
    return jsonify({
        'id': scenario.id,
        'name': scenario.name,
        'description': scenario.description,
        'type': scenario.type,
        'db_connection_id': scenario.db_connection_id,
        'created_at': scenario.created_at.isoformat(),
        'updated_at': scenario.updated_at.isoformat()
    })

# 数据库连接相关API
@app.route('/api/connections', methods=['GET'])
def get_connections():
    connections = DatabaseConnection.query.all()
    return jsonify([{
        'id': c.id,
        'name': c.name,
        'db_type': c.db_type,
        'host': c.host,
        'port': c.port,
        'database': c.database,
        'username': c.username,
        'created_at': c.created_at.isoformat()
    } for c in connections])

@app.route('/api/connections', methods=['POST'])
def create_connection():
    data = request.json
    connection = DatabaseConnection(
        name=data['name'],
        db_type=data['db_type'],
        host=data['host'],
        port=data['port'],
        database=data['database'],
        username=data['username'],
        password=data['password'],
        created_at=datetime.now()
    )
    db.session.add(connection)
    db.session.commit()
    return jsonify({'id': connection.id, 'message': 'Connection created successfully'}), 201

@app.route('/api/connections/<int:conn_id>/test', methods=['POST'])
def test_connection(conn_id):
    conn = DatabaseConnection.query.get_or_404(conn_id)
    try:
        db_connector = DatabaseConnector(
            db_type=conn.db_type,
            host=conn.host,
            port=conn.port,
            database=conn.database,
            username=conn.username,
            password=conn.password
        )
        success = db_connector.test_connection()
        return jsonify({'success': success})
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500

# SQL查询相关API
@app.route('/api/query', methods=['POST'])
def execute_query():
    data = request.json
    conn_id = data['connection_id']
    query = data['query']
    
    conn = DatabaseConnection.query.get_or_404(conn_id)
    
    try:
        # 执行查询
        db_connector = DatabaseConnector(
            db_type=conn.db_type,
            host=conn.host,
            port=conn.port,
            database=conn.database,
            username=conn.username,
            password=conn.password
        )
        
        result, columns = db_connector.execute_query(query)
        
        # 保存查询历史
        history = QueryHistory(
            connection_id=conn_id,
            query=query,
            executed_at=datetime.now(),
            success=True
        )
        db.session.add(history)
        db.session.commit()
        
        return jsonify({
            'success': True,
            'columns': columns,
            'data': result,
            'history_id': history.id
        })
    except Exception as e:
        # 保存失败的查询历史
        history = QueryHistory(
            connection_id=conn_id,
            query=query,
            executed_at=datetime.now(),
            success=False,
            error=str(e)
        )
        db.session.add(history)
        db.session.commit()
        
        return jsonify({'success': False, 'error': str(e)}), 500

# 初始化数据库表
with app.app_context():
    db.create_all()

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)
    